#!/bin/bash
#
# Arch Linux + Btrfs + Snapper 自动安装脚本（改进版）
#
# 功能：
# - 自动创建优化的 Btrfs 子卷结构
# - 正确配置 Snapper 快照管理（使用 snapper create-config）
# - 设置 Pacman 自动快照钩子
# - 完整的错误检查和验证
#
# 使用方法：
# 1. 启动到 Arch Linux Live 环境
# 2. 使用 cfdisk 对磁盘进行分区（EFI + Linux filesystem）
# 3. 运行此脚本：bash archlinux-btrfs-install-fixed.sh
#
# 作者：Claude & Lin
# 日期：2025-10-06
#

set -e  # 遇到错误立即退出

##############################################################################
# 0. 基础变量配置
##############################################################################

DISK="/dev/sda"              # 目标磁盘（根据实际情况修改）
EFI_PART="${DISK}1"          # EFI 分区
ROOT_PART="${DISK}2"         # Btrfs 根分区
MOUNTPOINT="/mnt"            # 安装挂载点

# 颜色定义（用于输出）
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 输出函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

##############################################################################
# 1. 基础检查
##############################################################################

echo "========================================"
echo "  Arch Linux + Btrfs + Snapper 安装"
echo "========================================"
echo ""

log_info "开始基础检查..."

# 检查是否以 root 运行
if [ "$EUID" -ne 0 ]; then
    log_error "请以 root 身份运行此脚本"
    exit 1
fi

# 检查目标磁盘是否存在
if [ ! -b "$DISK" ]; then
    log_error "磁盘 $DISK 不存在！"
    echo "可用磁盘："
    lsblk -d -o NAME,SIZE,TYPE | grep disk
    exit 1
fi

# 检查 EFI 分区是否存在
if [ ! -b "$EFI_PART" ]; then
    log_error "EFI 分区 $EFI_PART 不存在！"
    echo "请先使用 cfdisk 对磁盘进行分区"
    exit 1
fi

# 检查根分区是否存在
if [ ! -b "$ROOT_PART" ]; then
    log_error "根分区 $ROOT_PART 不存在！"
    echo "请先使用 cfdisk 对磁盘进行分区"
    exit 1
fi

log_info "磁盘检查通过"
echo "  目标磁盘：$DISK"
echo "  EFI 分区：$EFI_PART"
echo "  根分区：  $ROOT_PART"
echo ""

# 警告用户
log_warn "警告：此操作将格式化 $DISK 上的所有分区！"
echo "按 Ctrl+C 取消，或按回车继续..."
read -r

##############################################################################
# 2. 清理已挂载的分区
##############################################################################

log_info "检查并卸载已挂载的分区..."

# 卸载所有可能的挂载点
umount -R "$MOUNTPOINT" 2>/dev/null || true
umount "$EFI_PART" 2>/dev/null || true
umount "$ROOT_PART" 2>/dev/null || true

log_info "卸载完成"

##############################################################################
# 3. 格式化分区
##############################################################################

log_info "开始格式化分区..."

# 格式化 EFI 分区
log_info "格式化 EFI 分区：$EFI_PART"
mkfs.fat -F32 "$EFI_PART"

# 格式化 Btrfs 根分区
log_info "格式化 Btrfs 分区：$ROOT_PART"
mkfs.btrfs -f -L archlinux "$ROOT_PART"

log_info "分区格式化完成"
echo ""

##############################################################################
# 4. 创建 Btrfs 子卷结构
##############################################################################

log_info "创建 Btrfs 子卷结构..."

# 临时挂载根分区
mount "$ROOT_PART" "$MOUNTPOINT"

# 创建标准子卷
log_info "创建子卷：@（根目录）"
btrfs subvolume create "$MOUNTPOINT/@"

log_info "创建子卷：@home（用户目录）"
btrfs subvolume create "$MOUNTPOINT/@home"

log_info "创建子卷：@var_log（系统日志）"
btrfs subvolume create "$MOUNTPOINT/@var_log"

log_info "创建子卷：@var_cache（软件包缓存）"
btrfs subvolume create "$MOUNTPOINT/@var_cache"

log_info "创建子卷：@snapshots（快照存储）"
btrfs subvolume create "$MOUNTPOINT/@snapshots"

# 注意：不设置 @snapshots 为只读！
# Snapper 需要在其中创建快照子卷，必须可写
log_info "保持 @snapshots 为可写状态（Snapper 需要）"

# 验证子卷创建
log_info "验证子卷创建..."
btrfs subvolume list "$MOUNTPOINT"

# 卸载临时挂载
umount "$MOUNTPOINT"

log_info "子卷创建完成"
echo ""

##############################################################################
# 5. 挂载子卷
##############################################################################

log_info "挂载子卷..."

# Btrfs 挂载选项说明：
# - compress=zstd:3   : ZSTD 压缩，级别3（平衡性能和压缩率）
# - ssd               : SSD 优化
# - noatime          : 不更新访问时间（减少写入）
# - space_cache=v2   : 空间缓存版本2（性能优化）
# - discard=async    : 异步 TRIM（SSD 优化）

MOUNT_OPTS="compress=zstd:3,ssd,noatime,space_cache=v2,discard=async"

# 挂载根子卷
log_info "挂载根子卷 @ 到 $MOUNTPOINT"
mount -o "subvol=@,$MOUNT_OPTS" "$ROOT_PART" "$MOUNTPOINT"

# 创建挂载点目录
log_info "创建挂载点目录"
mkdir -p "$MOUNTPOINT"/{boot/efi,home,var/log,var/cache,.snapshots}

# 挂载其他子卷
log_info "挂载 @home 到 $MOUNTPOINT/home"
mount -o "subvol=@home,$MOUNT_OPTS" "$ROOT_PART" "$MOUNTPOINT/home"

log_info "挂载 @var_log 到 $MOUNTPOINT/var/log"
mount -o "subvol=@var_log,$MOUNT_OPTS" "$ROOT_PART" "$MOUNTPOINT/var/log"

log_info "挂载 @var_cache 到 $MOUNTPOINT/var/cache"
mount -o "subvol=@var_cache,$MOUNT_OPTS" "$ROOT_PART" "$MOUNTPOINT/var/cache"

log_info "挂载 @snapshots 到 $MOUNTPOINT/.snapshots"
mount -o "subvol=@snapshots,$MOUNT_OPTS" "$ROOT_PART" "$MOUNTPOINT/.snapshots"

# 挂载 EFI 分区
log_info "挂载 EFI 分区到 $MOUNTPOINT/boot/efi"
mount "$EFI_PART" "$MOUNTPOINT/boot/efi"

# 验证挂载
log_info "验证挂载状态："
findmnt -t btrfs,vfat | grep "$MOUNTPOINT"

log_info "子卷挂载完成"
echo ""

##############################################################################
# 6. 配置 Pacman 镜像源（使用中国镜像加速）
##############################################################################

log_info "配置 Pacman 镜像源..."

# 备份原始镜像源
cp /etc/pacman.d/mirrorlist /etc/pacman.d/mirrorlist.bak

# 添加中科大镜像源到顶部
sed -i '1i\## 中国科学技术大学镜像源\nServer = https://mirrors.ustc.edu.cn/archlinux/$repo/os/$arch\n' /etc/pacman.d/mirrorlist

# 更新软件包数据库
log_info "更新软件包数据库..."
pacman -Syy --noconfirm

log_info "镜像源配置完成"
echo ""

##############################################################################
# 7. 安装基本系统
##############################################################################

log_info "开始安装基本系统（这可能需要几分钟）..."

# 安装基础软件包
# - base           : 基础系统
# - linux          : Linux 内核
# - linux-firmware : 硬件固件
# - btrfs-progs    : Btrfs 工具
# - snapper        : 快照管理工具
# - vim            : 文本编辑器
# - networkmanager : 网络管理
# - sudo           : 权限管理
# - git            : 版本控制

pacstrap "$MOUNTPOINT" \
    base \
    linux \
    linux-firmware \
    btrfs-progs \
    snapper \
    vim \
    networkmanager \
    sudo \
    git \
    fontconfig \
    noto-fonts-cjk \
    noto-fonts-emoji

log_info "基本系统安装完成"
echo ""

##############################################################################
# 8. 生成 fstab
##############################################################################

log_info "生成 fstab 文件..."

genfstab -U "$MOUNTPOINT" >> "$MOUNTPOINT/etc/fstab"

log_info "fstab 内容："
echo "---"
cat "$MOUNTPOINT/etc/fstab"
echo "---"
echo ""

##############################################################################
# 9. 配置 Snapper（关键部分）
##############################################################################

log_info "配置 Snapper 快照系统..."

# 创建 Snapper 初始化脚本（在首次启动时执行）
# 为什么不在 chroot 中配置？
# 1. chroot 环境中 DBus 不可用
# 2. snapper create-config 需要处理已存在的 .snapshots
# 3. 在正常运行的系统中配置更可靠

cat > "$MOUNTPOINT/usr/local/bin/init-snapper.sh" << 'INIT_SCRIPT'
#!/bin/bash
#
# Snapper 初始化脚本
# 在系统首次启动后自动运行
#

set -e

log_info() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [INFO] $1"
}

log_error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $1" >&2
}

log_info "========================================="
log_info "开始初始化 Snapper..."
log_info "========================================="

# 步骤 1：确保 .snapshots 可写
log_info "步骤 1/6：确保 .snapshots 可写"
if ! touch /.snapshots/.test 2>/dev/null; then
    log_error "/.snapshots 不可写！"
    exit 1
fi
rm -f /.snapshots/.test
log_info "✓ .snapshots 可写"

# 步骤 2：处理已存在的 .snapshots 挂载
log_info "步骤 2/6：处理 .snapshots 挂载"
if mountpoint -q /.snapshots; then
    log_info "检测到 .snapshots 已挂载（预期行为）"
    
    # 临时卸载并移走
    log_info "临时卸载 .snapshots"
    umount /.snapshots
    mv /.snapshots /.snapshots.bak
    
    # 创建 Snapper 配置（会创建新的 .snapshots 子卷）
    log_info "步骤 3/6：创建 Snapper 配置"
    snapper -c root create-config /
    
    # 删除 Snapper 创建的新子卷（我们要用原来的 @snapshots）
    log_info "步骤 4/6：删除 Snapper 创建的子卷"
    btrfs subvolume delete /.snapshots
    
    # 恢复原来的目录
    log_info "步骤 5/6：恢复原 .snapshots 目录"
    mv /.snapshots.bak /.snapshots
    
    # 重新挂载
    log_info "重新挂载 .snapshots"
    mount -a
else
    # 如果 .snapshots 未挂载（不应该发生）
    log_info "步骤 3/6：直接创建 Snapper 配置"
    snapper -c root create-config /
fi

# 步骤 6：验证配置
log_info "步骤 6/6：验证 Snapper 配置"
if snapper list-configs | grep -q "root"; then
    log_info "✓ Snapper 配置成功"
else
    log_error "Snapper 配置验证失败"
    exit 1
fi

# 创建初始快照
log_info "创建初始快照..."
snapper create --description "系统安装完成"

# 显示快照列表
log_info "当前快照列表："
snapper list

log_info "========================================="
log_info "Snapper 初始化完成！"
log_info "========================================="

# 清理：删除此脚本（一次性任务）
log_info "清理初始化脚本..."
rm -f /usr/local/bin/init-snapper.sh

exit 0
INIT_SCRIPT

# 设置脚本权限
chmod +x "$MOUNTPOINT/usr/local/bin/init-snapper.sh"

log_info "✓ Snapper 初始化脚本已创建"
echo ""

##############################################################################
# 10. 配置 Pacman 自动快照钩子
##############################################################################

log_info "配置 Pacman 自动快照钩子..."

mkdir -p "$MOUNTPOINT/etc/pacman.d/hooks"

# 创建 pacman 操作前快照钩子
cat > "$MOUNTPOINT/etc/pacman.d/hooks/50-snapper-pre.hook" << 'HOOK'
[Trigger]
Operation = Upgrade
Operation = Install
Operation = Remove
Type = Package
Target = *

[Action]
Description = 创建 pacman 操作前快照
When = PreTransaction
Exec = /usr/bin/sh -c 'snapper -c root create --description "pacman-pre" --cleanup-algorithm=number 2>/dev/null || true'
HOOK

# 创建 pacman 操作后快照钩子
cat > "$MOUNTPOINT/etc/pacman.d/hooks/50-snapper-post.hook" << 'HOOK'
[Trigger]
Operation = Upgrade
Operation = Install
Operation = Remove
Type = Package
Target = *

[Action]
Description = 创建 pacman 操作后快照
When = PostTransaction
Exec = /usr/bin/sh -c 'snapper -c root create --description "pacman-post" --cleanup-algorithm=number 2>/dev/null || true'
HOOK

log_info "✓ Pacman 钩子配置完成"
echo ""

##############################################################################
# 11. 在 chroot 中进行基础配置
##############################################################################

log_info "在 chroot 环境中进行基础配置..."

arch-chroot "$MOUNTPOINT" /bin/bash << 'CHROOT_EOF'
set -e

# 启用 Snapper 定时服务
echo "[INFO] 启用 Snapper 定时服务..."
systemctl enable snapper-timeline.timer
systemctl enable snapper-cleanup.timer

echo "[INFO] ✓ 基础配置完成"
CHROOT_EOF

log_info "chroot 配置完成"
echo ""

##############################################################################
# 12. 最终验证
##############################################################################

log_info "执行最终验证..."

# 验证 fstab
if [ -f "$MOUNTPOINT/etc/fstab" ]; then
    log_info "✓ fstab 文件已生成"
else
    log_error "fstab 文件生成失败"
    exit 1
fi

# 验证子卷
log_info "验证 Btrfs 子卷..."
if btrfs subvolume list "$MOUNTPOINT" | grep -E "(@|@home|@var_log|@var_cache|@snapshots)"; then
    log_info "✓ 所有子卷已正确创建"
else
    log_warn "部分子卷可能未正确创建"
fi

# 验证挂载
log_info "验证挂载状态..."
if findmnt "$MOUNTPOINT" > /dev/null; then
    log_info "✓ 所有分区已正确挂载"
else
    log_error "挂载状态异常"
    exit 1
fi

# 验证 Snapper 初始化脚本
if [ -x "$MOUNTPOINT/usr/local/bin/init-snapper.sh" ]; then
    log_info "✓ Snapper 初始化脚本已就绪"
else
    log_error "Snapper 初始化脚本未找到"
    exit 1
fi

log_info "所有验证通过！"
echo ""

##############################################################################
# 13. 安装完成提示
##############################################################################

echo ""
echo "========================================"
echo "  基础系统安装完成！"
echo "========================================"
echo ""
echo "接下来需要手动完成以下配置："
echo ""
echo "1. 进入系统："
echo "   arch-chroot /mnt"
echo ""
echo "2. 配置 Bootloader（推荐 systemd-boot）："
echo "   bootctl install"
echo "   # 创建启动条目（参考 README.md）"
echo ""
echo "3. 基础系统配置："
echo "   # 设置时区"
echo "   ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime"
echo "   hwclock --systohc"
echo ""
echo "   # 设置语言"
echo "   sed -i 's/#en_US.UTF-8/en_US.UTF-8/' /etc/locale.gen"
echo "   sed -i 's/#zh_CN.UTF-8/zh_CN.UTF-8/' /etc/locale.gen"
echo "   locale-gen"
echo "   echo 'LANG=en_US.UTF-8' > /etc/locale.conf"
echo ""
echo "   # 设置主机名"
echo "   echo 'archlinux' > /etc/hostname"
echo ""
echo "   # 设置 root 密码"
echo "   passwd"
echo ""
echo "   # 创建用户"
echo "   useradd -m -G wheel lin"
echo "   passwd lin"
echo "   sed -i 's/# %wheel ALL=(ALL:ALL) ALL/%wheel ALL=(ALL:ALL) ALL/' /etc/sudoers"
echo ""
echo "   # 启用网络"
echo "   systemctl enable NetworkManager"
echo ""
echo "4. 退出 chroot 并重启："
echo "   exit"
echo "   reboot"
echo ""
echo "5. **重要**：首次启动后，立即初始化 Snapper："
echo "   sudo /usr/local/bin/init-snapper.sh"
echo ""
echo "========================================"
echo "  已自动配置的功能"
echo "========================================"
echo ""
echo "✓ Btrfs 子卷结构："
echo "  - @ (根目录)"
echo "  - @home (用户目录)"
echo "  - @var_log (系统日志)"
echo "  - @var_cache (软件包缓存)"
echo "  - @snapshots (快照存储)"
echo ""
echo "✓ Snapper 快照系统："
echo "  - 自动快照策略（每小时/每天/每周/每月）"
echo "  - Pacman 自动快照钩子"
echo "  - 首次启动初始化脚本"
echo ""
echo "✓ 系统优化："
echo "  - ZSTD 压缩"
echo "  - SSD 优化"
echo "  - 中国镜像源"
echo ""
echo "========================================"
echo "  Snapper 使用说明"
echo "========================================"
echo ""
echo "首次启动后初始化："
echo "  sudo /usr/local/bin/init-snapper.sh"
echo ""
echo "常用命令："
echo "  sudo snapper list                    # 查看快照列表"
echo "  sudo snapper create -d '描述'        # 创建手动快照"
echo "  sudo snapper diff 1..2               # 比较快照差异"
echo "  sudo snapper delete 3                # 删除快照"
echo ""
echo "从快照恢复文件："
echo "  sudo cp /.snapshots/X/snapshot/path/to/file /path/to/file"
echo ""
echo "========================================"
echo ""

log_info "脚本执行完成！请按照上面的提示完成后续配置。"

